#pragma once

#include "DnsLayer.h"
#include "DnsLayerEnums.h"
#include "DnsResourceData.h"
#include <stdio.h>
#include <string>
#include <stdint.h>

/// @file

/// @namespace pcpp
/// @brief The main namespace for the PcapPlusPlus lib
namespace pcpp
{
	// forward declarations
	class DnsLayer;
	class IDnsResourceData;
	class DnsResourceDataPtr;

	/// @class IDnsResource
	/// An abstract class for representing all types of DNS records. This class gives access to all available record
	/// data such as DNS type, class, name, type of record, etc. The DnsLayer holds an instance of (inherited type of)
	/// this class for each DNS record in the DNS packet
	class IDnsResource
	{
	protected:
		friend class DnsLayer;
		friend class IDnsResourceData;

	protected:
		DnsLayer* m_DnsLayer;
		size_t m_OffsetInLayer;
		IDnsResource* m_NextResource;
		std::string m_DecodedName;
		size_t m_NameLength;
		uint8_t* m_ExternalRawData;

		IDnsResource(DnsLayer* dnsLayer, size_t offsetInLayer);

		IDnsResource(uint8_t* emptyRawData);

		size_t decodeName(const char* encodedName, char* result, int iteration = 1);
		void encodeName(const std::string& decodedName, char* result, size_t& resultLen);

		IDnsResource* getNextResource() const
		{
			return m_NextResource;
		}
		void setNextResource(IDnsResource* next)
		{
			m_NextResource = next;
		}

		uint8_t* getRawData() const;

		void setDnsLayer(DnsLayer* dnsLayer, size_t offsetInLayer);

	public:
		virtual ~IDnsResource() = default;

		/// @return The DNS type of this record
		DnsType getDnsType() const;

		/// Set DNS type for this record
		/// @param[in] newType The type to set
		void setDnsType(DnsType newType);

		/// @return The DNS class of this record
		DnsClass getDnsClass() const;

		/// Set DNS class for this record
		/// @param[in] newClass The class to set
		void setDnsClass(DnsClass newClass);

		/// @return The name of this record
		const std::string& getName() const
		{
			return m_DecodedName;
		}

		/// @return The record name's offset in the packet
		size_t getNameOffset() const
		{
			return m_OffsetInLayer;
		}

		/// Set the name of this record. The input name can be a standard hostname (e.g 'google.com'), or it may contain
		/// a pointer to another string in the packet (as explained here: http://www.zytrax.com/books/dns/ch15/#name).
		/// The pointer is used to reduce the DNS packet size and avoid unnecessary duplications. In case you
		/// want to use a pointer in your string you should use the following format: 'some.domain.#{offset}' where
		/// '#{offset}' is a the offset from the start of the layer. For example: if the string 'yahoo.com' already
		/// appears in offset 12 in the packet and you want to set the name of the current record to
		/// 'my.subdomain.yahoo.com' you may use the following string: 'my.subdomain.#12'. This will result in writing
		/// 'my.subdomain' and a pointer to offset 12.<BR> Please notice the new name can be shorter or longer of the
		/// old name, so this method can cause the packet to be shorten or extended
		/// @param[in] newName The name to set
		/// @return True if name was set successfully or false if input string is malformed or if an error occurred
		bool setName(const std::string& newName);

		// abstract methods

		/// @return The total size in bytes of this record
		virtual size_t getSize() const = 0;

		/// @return The type of this record (query, answer, authority, additional)
		virtual DnsResourceType getType() const = 0;
	};

	/// @class DnsQuery
	/// Representing a DNS query record
	class DnsQuery : public IDnsResource
	{
		friend class DnsLayer;

	private:
		DnsQuery(DnsLayer* dnsLayer, size_t offsetInLayer) : IDnsResource(dnsLayer, offsetInLayer)
		{}

		explicit DnsQuery(uint8_t* emptyRawData) : IDnsResource(emptyRawData)
		{}

	public:
		~DnsQuery() override = default;

		// implementation of abstract methods
		size_t getSize() const override
		{
			return m_NameLength + 2 * sizeof(uint16_t);
		}
		DnsResourceType getType() const override
		{
			return DnsQueryType;
		}
	};

	/// @class DnsResource
	/// Representing DNS record other than DNS query
	class DnsResource : public IDnsResource
	{
		friend class DnsLayer;

	private:
		DnsResourceType m_ResourceType;

		DnsResource(DnsLayer* dnsLayer, size_t offsetInLayer, DnsResourceType resourceType)
		    : IDnsResource(dnsLayer, offsetInLayer)
		{
			m_ResourceType = resourceType;
		}

		DnsResource(uint8_t* emptyRawData, DnsResourceType resType)
		    : IDnsResource(emptyRawData), m_ResourceType(resType)
		{}

	public:
		~DnsResource() override = default;

		/// @return The time-to-leave value for this record
		uint32_t getTTL() const;

		/// Set time-to-leave value for this record
		/// @param[in] newTTL The new TTL value to set
		void setTTL(uint32_t newTTL);

		/// @return The data length value for this record (taken from the "data length" field of the record)
		size_t getDataLength() const;

		/// @return A smart pointer to an IDnsResourceData object that contains the DNS resource data. It is guaranteed
		/// that the smart pointer will always point to an object and never to nullptr. The specific object type depends
		/// on the DNS type of this record:<BR>
		/// - For type A (::DNS_TYPE_A): the return value is a smart pointer to IPv4DnsResourceData object that contains
		/// the IPv4 address<BR>
		/// - For type AAAA (::DNS_TYPE_AAAA): the return value is a smart pointer to IPv6DnsResourceData object that
		/// contains the IPv6 address<BR>
		/// - For types NS, CNAME, DNAME, PTR (::DNS_TYPE_NS, ::DNS_TYPE_CNAME, ::DNS_TYPE_DNAM, ::DNS_TYPE_PTR): the
		/// return value is a smart pointer to StringDnsResourceData object that contains the name<BR>
		/// - For type MX (::DNS_TYPE_MX): the return value is a smart pointer to MxDnsResourceData object that contains
		/// the MX data (preference and mail exchange name)<BR>
		/// - For all other types: the return value is a smart pointer to GenericDnsResourceData which contains a byte
		/// array of the data
		DnsResourceDataPtr getData() const;

		/// @return The offset of data in the DNS layer
		size_t getDataOffset() const;

		/// Set resource data. The given IDnsResourceData input object is validated against the DNS type of the
		/// resource. For example: if DNS type is A and data isn't of type IPv4DnsResourceData (which contains the IPv4
		/// address) a log error will be printed and the method will return false. This method currently supports the
		/// following DNS types:<BR>
		/// - ::DNS_TYPE_A (IPv4 address) - data is expected to be a pointer to IPv4DnsResourceData with a valid IPv4
		/// address
		/// - ::DNS_TYPE_AAAA (IPv6 address) - data is expected to be a pointer to IPv6DnsResourceData with a valid IPv6
		/// address
		/// - ::DNS_TYPE_NS, ::DNS_TYPE_CNAME, ::DNS_TYPE_DNAM, ::DNS_TYPE_PTR (name data) - data is expected to be a
		/// pointer to StringDnsResourceData object that contains a host name, e.g: 'www.google.com'
		/// - ::DNS_TYPE_MX (MX data) - data is expected to be a pointer to MxDnsResourceData object that contains the
		/// MX data
		/// - else: data is expected to be a pointer to GenericDnsResourceData object that contains a valid hex string
		/// (valid hex string means a string which has an even number of characters representing a valid hex data. e.g:
		/// '0d0a45569a9b')
		/// @param[in] data The pointer to the data object, as described above
		/// @return True if data was properly set or false if data is illegal or method couldn't extend or shorted the
		/// packet (appropriate error log is printed in all cases)
		bool setData(IDnsResourceData* data);

		/// Some records don't have a DNS class and the bytes used for storing the DNS class are used for other purpose.
		/// This method enables the user to receive these bytes
		/// @return The value stored in this place
		uint16_t getCustomDnsClass() const;

		/// Some records don't have a DNS class and the bytes used for storing the DNS class are used for other purpose.
		/// This method enables the user to set these bytes
		/// @param[in] customValue The value to set
		void setCustomDnsClass(uint16_t customValue);

		// implementation of abstract methods
		size_t getSize() const override
		{
			return m_NameLength + 3 * sizeof(uint16_t) + sizeof(uint32_t) + getDataLength();
		}
		DnsResourceType getType() const override
		{
			return m_ResourceType;
		}
	};

}  // namespace pcpp
